package list

import (
	bh "github.com/emirpasic/gods/trees/binaryheap"
	"github.com/emirpasic/gods/utils"
)

// Find median from data stream
// Problem Definition: https://leetcode.com/problems/find-median-from-data-stream/

type MedianFinder struct {
	large, small *bh.Heap
}

func largePriority(a, b interface{}) int {
	priorityA := a.(int)
	priorityB := b.(int)
	return -utils.IntComparator(priorityA, priorityB)
}

func minPriority(a, b interface{}) int {
	priorityA := a.(int)
	priorityB := b.(int)
	return utils.IntComparator(priorityA, priorityB)
}

func Constructor() MedianFinder {
	large := bh.NewWith(largePriority)
	min := bh.NewWith(minPriority)

	return MedianFinder{large, min}
}

func (m *MedianFinder) FindMedian() float64 {
	s, _ := m.small.Peek()
	l, _ := m.large.Peek()
	if m.large.Size() < m.small.Size() {
		return float64(s.(int))
	} else if m.small.Size() < m.large.Size() {
		return float64(l.(int))
	}
	return float64(l.(int)+s.(int)) / 2.0
}

func (m *MedianFinder) AddNum(num int) {
	if m.small.Size() >= m.large.Size() {
		m.small.Push(num)
		val, _ := m.small.Pop()
		m.large.Push(val)
	} else {
		m.large.Push(num)
		val, _ := m.large.Pop()
		m.small.Push(val)
	}
}
